import torchvision
import torch
from ModelStatistics import *
class ResnetClass10(torch.nn.Module):
    def __init__(self):
        super(ResnetClass10, self).__init__()
        self.cnn_layers = torchvision.models.resnet50(pretrained=True)
        # Get the last layer from resnet50.
        num_ftrs = self.cnn_layers.fc.in_features
        # Modify the last layer output features dim to classes size(10).
        self.cnn_layers.fc = torch.nn.Linear(num_ftrs, 10)

    def forward(self, x):
        out = self.cnn_layers(x)
        return out

def test_model():
    model = ResnetClass10()
    all_param_size(model)
    all_param_name_and_size(model)
    print_model(model)

# test_model()